######################################################################
# Time-stamp: <2006-09-28 22:01:41 patrick>
#
# Package: EP::Core::Data::Parser::seq
#
# Expression Data Parser Classes
######################################################################
package EP::Core::Data::Parser::seq;

use strict;
use base qw ( EP::Core::Data::Parser );

use EP::Config;
use EP::Common;
use EP::Common::General;
use EP::Common::Error qw ( :try );
use EP::EPC::Core;
use EP::EPC::Core::Component;

use BioMart::Initializer;
use BioMart::AttributeTable;
use BioMart::Query;
use BioMart::QueryPlanner;
use BioMart::ResultTable;
use Bio::Seq;
use Bio::DB::GFF;

use File::Temp;
use File::Spec;

sub parse_data {
  my $self = shift;

  my $data = $self->upload_data ();

  my $cur_folder = $self->{epc}->{session}->current_folder;
  if ( $self->_process_uploaded_data ( $data ) ) {
    my ($dataset) = EP::Common::DBI->do_transaction (
      sub {
	my ($self) = @_;
	my $dataset = EP::Common::DBI::Dataset->create ({ name => $data->{dataset_name} . " : ". localtime,
							  description  => $self->{query}->param ( "dataset_description" ) . "",
							  type         => "seq",
							  species      => $self->{query}->param ( "dataset_species" ) . "",
							  folder       => $cur_folder->id,
							  ep_session      => $self->{epc}->{current_user}->session_only ? $self->{epc}->{session}->id : undef,
							  owner        => $self->{epc}->{current_user},
							  group_write  => 0,
							  group_read   => 1,
							  other_read   => 0
							});
	$self->_write_dataset_data ( $dataset, $data ) or die "Failed to write the dataset data!";
	return $dataset;
      }, $self );

    return $dataset;
  }

  die "Failed to process uploaded data.";
}

sub _process_uploaded_data {
  my $self = shift;
  my ($data) = @_;

  # Bit of a hack.... BioSeq and File::Temp bite eachother because of overloading....
  my $file = "$data->{data}";

  $data->{sequences} = Bio::SeqIO->new( -file  => $file,
					-format => $self->{query}->param('seq_format') );

  unlink $data->{data};

  return 1;
}

# sub: _get_sequences
#   Sequence retrieval routine. It will connect to either a BioMart or GFF database instance, depending
#   on what has been defined in the configuration and retrieve either all
#   sequences for a certain species of type 'gene' or those sequences for which their ids where given.
# parameters:
#   species - Which species to obtain the sequences from
#   ids - Optional. Reference to an array containing the ids of the sequences to be obtained
# returns:
#   A sequence stream object, either an BioMart::ResultTable or Bio::SeqIO object
# see also:
#   <BioMart::ResultTable> <Bio::SeqIO>
sub _get_sequences {
    my $self = shift;
    my ($species, $ids, $upstream_bp) = @_;
    my ($stream);

    my $BioMart = $EP::Config::h->{BioMart}; # get global BioMart configuration
    my $BioMartSp = $EP::Config::h->{SPECIES_METADATA}{$species}{BioMart}; # get species specific BioMart configuration

    if ( defined( $BioMartSp ) ) {
	my $init;
	# a BioMart sequenceDataset has been defined, obtain sequences from BioMart
	warn "Retrieving sequences from BioMart instance" if DEBUG;

	# setup BioMart, do this inside an eval as Biomart throws ugly tracestacks as errors and dies
	eval {
	    $init = BioMart::Initializer->new( 'registryFile' => $BioMart->{confFile} );
	    };
	if ( $@ ) {
	    warn $@ if DEBUG;
	    throw EP::Common::Error ( -text => "Unable to connect to a BioMart instance to retrieve the sequences, please try again later" );
	}

	my $registry = $init->getRegistry;

	$BioMart->{ensConfTree} = $registry->getDatasetByName( $BioMart->{virtualSchema}, 
							       $BioMartSp->{ensDataset} )->getConfigurationTree();
	$BioMart->{structConfTree} = $registry->getDatasetByName( $BioMart->{virtualSchema}, 
								  $BioMartSp->{structDataset} )->getConfigurationTree();
	$BioMart->{seqConfTree} = $registry->getDatasetByName( $BioMart->{virtualSchema}, 
							       $BioMartSp->{seqDataset} )->getConfigurationTree();
	$registry->createAllLinks();
	
	# setup BioMart query
	my $query = BioMart::Query->new( 'registry' => $registry,
					 'virtualSchemaName' => $BioMart->{virtualSchema} );
#	$query->virtualSchema( $BioMart->{virtualSchema} );
	$self->_BioMart_add_seq_type( $BioMart, \$query);
	
	# only retrieve upstream sequences with a certain length
	if ( $upstream_bp ) {
		$self->_BioMart_add_filter( \$query, $BioMart->{seqConfTree}, 'upstream_flank', $upstream_bp );
	} else {
	    $self->_BioMart_add_filter( \$query, $BioMart->{seqConfTree}, 'coding' );
	}

	if ( $ids ) {
	    # limit search to certain identifiers
	    $self->_BioMart_add_list_filter( \$query, $BioMart->{ensConfTree}, $BioMartSp->{xrefType} , $ids );
	}

	# add links
	my $path = $registry->getPath( $BioMart->{virtualSchema}, $BioMartSp->{ensDataset}, $BioMartSp->{seqDataset} );
	
	my $firstDataset = shift @{ $path };
	my $lastDataset;
	my $alternate = $BioMart->{seqType};
	foreach my $thisDataset ( reverse( @{ $path } ) ) {
	    if ( $lastDataset ) {
		# use the seq type as the alternate link between the next to last, and the last
		# dataset only
		$self->_BioMart_addLinks( $BioMart, \$query, $registry, $thisDataset, $lastDataset, $alternate );
		$alternate = undef;
	    }
	    $lastDataset = $thisDataset;
	}
	$self->_BioMart_addLinks( $BioMart, \$query, $registry, $firstDataset, $lastDataset, $alternate );
	
	my $queryPlanner = BioMart::QueryPlanner->new();
	$stream = $queryPlanner->getResultTable( $query, $registry, undef, $BioMart->{initBatchSize} );
	
	return $stream;
    } else {
	# no way to retrieve sequences for this organism
	throw EP::Common::Error ( -text => "Genome sequence not available for species $species" );
    }
}

# sub: _write_dataset_data
#   Writes the sequence data from the dataset to a file on disk in fasta format
#   Expects the sequence data to be an instance of a <BioMart::ResultTable> or <Bio::SeqIO> object
#   If the start and end basepair position are set for the dataset object, 
#   the corresponding subpart of the sequences is written to file
# parameters:
#   dataset - Which dataset to write to disk
# returns:
#   1 upon successfull writing of the sequences to the file
# see also:
#   <EP::Core::Dataset::Data::seq::start_bp> <EP::Core::Dataset::Data::seq::end_bp>
sub _write_dataset_data {
    my $self = shift;
    my ($dataset, $data) = @_;
    my (@seq_notfound);

    my $stream = defined $data ? $data->{sequences} : $dataset->data->{sequences};
    my $folder   = $EP::Config::h->{EPNG_USRDATA_PATH} . "/" . $dataset->folder->location;
    my $filename = $dataset->filename;
    my $start_bp = $dataset->data->start_bp();
    my $end_bp = $dataset->data->end_bp();

    # always write fasta format
    my $out = Bio::SeqIO->new ( -file => ">$folder/$filename.fasta",
				-format => "fasta" );

    if ( ref( $stream ) =~ /^BioMart::ResultTable/ ) {
	# sequences obtained from BioMart
	unless( $stream->hasMoreRows ) {
	    throw EP::Common::Error ( -text => "No sequences have been found" );
	}

 	while ( $stream->hasMoreRows ) {
 	    my $row = $stream->nextRow;
	    my $seq = Bio::Seq->new( -seq => pop @{ $row },
				     -id => join( '|', @{ $row } ) . ":$start_bp,$end_bp"
				     );
	    $out->write_seq( $seq )
		or die "Unable to write sequence " . $seq->display_id . " to file $filename.fasta\n";
	}
    } elsif ( ref( $stream ) =~ /^Bio::DB::GFF/ ) {
        # sequences obtained from GFF

   	while ( 1 ) {
	    my $seq = 1;
	    my ($sequence, $id);

	    ### UGlY HACK ###
	    # the get_seq_stream and get_Stream_by_name do not actually return the same objects
	    # so we can only get the id if it is coming from get_Stream_by_name
	    my $try_id = $stream->{ids}->[0] if ref( $stream ) eq 'Bio::DB::GFF::ID_Iterator';

	    try {
		$seq = $stream->next_seq();
		
		if ( $start_bp && $end_bp ) {
		    # Do we only want part of the sequence ?
		    $sequence = $seq->subseq ( $start_bp, $end_bp )->seq;
		    $id = $seq->display_id . ":$start_bp,$end_bp";
		} else {
		    $sequence = $seq->seq;
		    $id = $seq->display_id;
		}

		# Need to recreate a new Bio::Seq object, as the get_seq_stream object resets
		# the display_id to the chromosome sequence if a subsequence is selected.
		my $seqObj = Bio::Seq->new( -seq => $sequence,
					    -id => $id
					    );
		$out->write_seq($seqObj) or
		    die "Unable to write sequence $id to file $filename.fasta\n";

	    } otherwise {
		push @seq_notfound, $try_id if $try_id;
	    };
	    last if not $seq;
	}
	
	# print warning if some id's where not found
	warn "The sequences for the following IDs were not obtained: " . join( " ", @seq_notfound) 
	    if scalar(@seq_notfound) > 0;
    } elsif ( ref ( $stream ) =~ /^Bio::Seq/ ) {
      warn "Writing uploaded data";
      my $oseq;
      try {
	while ( $oseq = $stream->next_seq ) {
	  $out->write_seq ( $oseq );
	}
      } otherwise {
	die "There was a problem writing the uploaded data.\n";
      }
    } else {
	# don't know how to deal with these sequences
	die "Do not know how to obtain the sequences from this type of stream: " . ref( $stream );
    }
    
    `perl -i -ne 'chomp if not /^>/; print "\n" if /^>/ and \$. > 1; print;' $folder/$filename.fasta`;
    warn "Successfully wrote sequence dataset $filename.fasta to folder $folder" if DEBUG;
    
    return 1;
}

# sub: _BioMart_add_seq_type
#   Adds a certan sequences type to a preexisting BioMart query.
#   It will also check whether this sequence type is valid, and if not throw an error
# parameters:
#   BioMart - a hash table containing all relevant information about the BioMart instance
#   queryref - reference to the current BioMart query
# returns:
#   null
# see also:
#   <BioMart::Initializer> <BioMart::Query>
sub _BioMart_add_seq_type {
    my $self = shift;
    my ( $BioMart, $queryref ) = @_;

    my $att = $BioMart->{seqConfTree}->getAttributeByName( $BioMart->{seqType} ) 
	or die "Sequence $BioMart->{seqType} not supported by [" . $BioMart->{seqConfTree}->dataSetName . "]\n";
    
    ${ $queryref }->addAttribute($att);

    # switch seqtype to add other attributes to simulate martview header
    $self->_BioMart_header_attributes( $BioMart, $queryref );
}

# sub: _BioMart_header_attributes
#   Retrieves all BioMart header attributes as defined in the BioMart configuration
# parameters:
#   BioMart - a hash table containing all relevant information about the BioMart instance
#   queryref - reference to the current BioMart query
# returns:
#   null
# see also:
#   <BioMart::Initializer>
sub _BioMart_header_attributes {
    my $self = shift;
    my ( $BioMart, $queryref ) = @_;

    my $seq_proto;
    if ($BioMart->{seqType} =~ m/(transcript\w*|gene\w*)$/ ) {
	$seq_proto = $1;
    } else {
	map { $seq_proto = $_ } split /\_/, $BioMart->{seqType}; #will resolve to the last element in the seq_type
    }

    foreach my $attname ( @{ $BioMart->{headerAttNames}{ $seq_proto } } ) {
	$self->_BioMart_add_attribute($BioMart, $queryref, $attname);
    }
}

# sub: _BioMart_add_attribute
#   Adds a certan attribute to a preexisting BioMart query.
#   It will also check whether this attribute is valid, and if not throw an error
# parameters:
#   BioMart - a hash table containing all relevant information about the BioMart instance
#   queryref - reference to the current BioMart query
#   name - attribute name
# returns:
#   null
# see also:
#   <BioMart::Initializer> <BioMart::Query>
sub _BioMart_add_attribute {
    my $self = shift;
    my ( $BioMart, $queryref, $name ) = @_;

    my $att = $BioMart->{structConfTree}->getAttributeByName( $name );

    unless ( $att ) {
	$att = $BioMart->{ensConfTree}->getAttributeByName( $name );
    }
    
    die "No Attribute $name in [" . $BioMart->{structConfTree}->dataSetName . "] or [" . $BioMart->{ensConfTree}->dataSetName . "]\n" unless ( $att );
    ${ $queryref }->addAttribute( $att );
}

# sub: _BioMart_add_filter
#   Adds a certan filter, belonging to a certain configurationTree (BioMart instance) to a preexisting BioMart query.
#   It will also check whether this filter is valid, and if not throw an error
# parameters:
#   queryref - reference to the current BioMart query
#   confTree - BioMart configurationTree
#   name - filter name
#   value - filter value
# returns:
#   null
# see also:
#   <BioMart::Initializer> <BioMart::Query>
sub _BioMart_add_filter {
    my $self = shift;
    my ( $queryref, $confTree, $name, $value ) = @_;
    die "No name provided for BioMart filter" if $name =~ /^\s*$/;

    my $filt = $confTree->getFilterByName( $name ); # first try filter

    unless( $filt ) {
	# then try option
	$filt = $confTree->getOptionByName( $name )->filter;
    }
    die "No Filter $name in [" . $confTree->dataSetName . "]\n" unless( $filt );

    my $atbl = BioMart::AttributeTable->new();
    $atbl->addRow( [ $value ] );
    $filt->setTable( $atbl );
    ${ $queryref }->addFilter( $filt );
}

# sub: _BioMart_add_list_filter
#   Adds a certan list of filter values, belonging to a certain configurationTree (BioMart instance) 
#   to a preexisting BioMart query.
#   It will also check whether these filter are valid, and if not throw an error
# parameters:
#   queryref - reference to the current BioMart query
#   confTree - BioMart configurationTree
#   name - filter name
#   ids - filter values
# returns:
#   null
# see also:
#   <BioMart::Initializer> <BioMart::Query>
sub _BioMart_add_list_filter {
    my $self = shift;
    my ($queryref, $confTree, $name, $ids ) =  @_;
    die "No name provided for BioMart filter" if $name =~ /^\s*$/;

    my $filt = $confTree->getFilterByName( $name ); # first try filter
    unless( $filt ) {
	# then try option
	my $opt = $confTree->getOptionByName( $name );
	if ( $opt ) {
	    $filt = $opt->filter;
	}
    }
    die "No Filter $name in [" . $confTree->dataSetName . "]\n" unless( $filt );

    my $atbl = BioMart::AttributeTable->new();
    foreach my $value ( @{ $ids } ) {
        $value =~ s/'/''/; # escape for SQL
	$atbl->addRow( [ $value ] );
    }
    
    $filt->setTable( $atbl );
    ${ $queryref }->addFilter( $filt );
}

# sub: _BioMart_addLinks
#   Will setup the links between the different datasets within the BioMart instance and add them to the 
#   current BioMart query
# parameters:
#   BioMart -  a hash table containing all relevant information about the BioMart instance
#   queryref - reference to the current BioMart query
#   registry - the initialized BioMart connection
#   source_ds - source dataset
#   target_ds - target dataset
#   alternate - alternative link if default fails
# returns:
#   null
# see also:
#   <BioMart::Initializer> <BioMart::Query>
sub _BioMart_addLinks {
    my $self = shift;
    my ( $BioMart, $queryref, $registry, $source_ds, $target_ds, $alternate ) = @_;
    
    my $link = $registry->getLinkBetween( $BioMart->{virtualSchema}, $source_ds, $target_ds )
	or die "No link between $source_ds and $target_ds\n";
    $link->defaultLink( $alternate ) if $alternate ;
    ${ $queryref }->addLinks( $link );
}

1;
